#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
更新用户表脚本
用于向已存在的users表添加所有缺失的字段
"""

# 导入系统相关模块
import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 导入dotenv用于加载环境变量
from dotenv import load_dotenv
# 导入SQLAlchemy相关组件
from sqlalchemy import text

# 加载环境变量
load_dotenv()

def check_and_update_users_table():
    """
    检查并更新users表结构
    向已存在的users表添加所有缺失的字段
    """
    print("开始检查和更新users表结构...")
    
    try:
        # 导入数据库引擎
        from utils.models import engine
        
        # 创建连接
        with engine.connect() as connection:
            # 开始事务
            trans = connection.begin()
            
            try:
                # 获取当前表结构
                result = connection.execute(text("SHOW COLUMNS FROM users"))
                existing_columns = [row[0] for row in result.fetchall()]
                print(f"当前users表字段: {existing_columns}")
                
                # 需要的字段列表
                required_columns = {
                    'role': "VARCHAR(20) DEFAULT 'researcher'",
                    'is_active': "INTEGER DEFAULT 1",
                    'created_at': "DATETIME DEFAULT CURRENT_TIMESTAMP",
                    'updated_at': "DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP"
                }
                
                # 检查并添加缺失的字段
                for column_name, column_definition in required_columns.items():
                    if column_name not in existing_columns:
                        sql = f"ALTER TABLE users ADD COLUMN {column_name} {column_definition}"
                        connection.execute(text(sql))
                        print(f"已添加字段: {column_name}")
                    else:
                        print(f"字段已存在: {column_name}")
                
                # 提交事务
                trans.commit()
                print("users表结构更新成功!")
                
            except Exception as e:
                # 回滚事务
                trans.rollback()
                raise e
                
    except Exception as e:
        print(f"更新users表结构时出错: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

# 程序入口
if __name__ == "__main__":
    check_and_update_users_table()